{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import safetensors\n",
    "import torch\n",
    "import json\n",
    "\n",
    "def check_cached_latent(path):\n",
    "    try:\n",
    "        with safetensors.safe_open(path, framework=\"pt\", device=\"cpu\") as f:\n",
    "            metadata = f.metadata()\n",
    "            keys = f.keys()\n",
    "            \n",
    "            print(f\"File: {os.path.basename(path)}\")\n",
    "            print(f\"Metadata present: {'Yes' if metadata else 'No'}\")\n",
    "            print(f\"Keys in the safetensor: {', '.join(keys)}\")\n",
    "            \n",
    "            if 'latent' in keys:\n",
    "                latent = f.get_tensor('latent')\n",
    "                print(f\"Latent shape: {latent.shape}\")\n",
    "                print(f\"Latent dtype: {latent.dtype}\")\n",
    "            \n",
    "            if metadata:\n",
    "                print(\"\\nMetadata:\")\n",
    "                for key, value in metadata.items():\n",
    "                    try:\n",
    "                        parsed_value = json.loads(value)\n",
    "                        print(f\"  {key}: {parsed_value}\")\n",
    "                    except json.JSONDecodeError:\n",
    "                        print(f\"  {key}: {value}\")\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(f\"Error processing the cached latent file: {str(e)}\")\n",
    "\n",
    "def scan_latent_cache(cache_dir):\n",
    "    print(f\"Scanning latent cache directory: {cache_dir}\")\n",
    "    for root, _, files in os.walk(cache_dir):\n",
    "        for file in files:\n",
    "            if file.endswith('.safetensors'):\n",
    "                full_path = os.path.join(root, file)\n",
    "                print(\"\\n\" + \"=\"*50)\n",
    "                check_cached_latent(full_path)\n",
    "\n",
    "#scan file instead of directory\n",
    "def scan_latent_file(file_path):\n",
    "    print(f\"Scanning latent file: {file_path}\")\n",
    "    check_cached_latent(file_path)\n",
    "\n",
    "# Example usage\n",
    "if __name__ == \"__main__\":\n",
    "    file_path = r\"P:\\downloads\\latents\\_latent_cache\\img_028_bLPQW46Jb_ve7lfwu4HXEA.safetensors\"\n",
    "    scan_latent_file(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from diffusers import AutoencoderTiny\n",
    "import safetensors\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def load_vae(device='cuda' if torch.cuda.is_available() else 'cpu'):\n",
    "    vae = AutoencoderTiny.from_pretrained(\"madebyollin/taef1\", torch_dtype=torch.bfloat16).to(device)\n",
    "    return vae\n",
    "\n",
    "def load_latent(path):\n",
    "    try:\n",
    "        with safetensors.safe_open(path, framework=\"pt\", device=\"cpu\") as f:\n",
    "            latent = f.get_tensor('latent')\n",
    "        return latent\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading latent from {path}: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "def decode_latents(latents: torch.Tensor, vae: AutoencoderTiny):\n",
    "    latents = latents.to(vae.device, dtype=torch.bfloat16)\n",
    "    latents = latents / vae.config.scaling_factor\n",
    "    with torch.no_grad():\n",
    "        images = vae.decode(latents).sample\n",
    "    return images\n",
    "\n",
    "def show_images(images, name='Decoded Images'):\n",
    "    plt.figure(figsize=(12, 12))\n",
    "    plt.suptitle(name, fontsize=16)\n",
    "    for i, img_tensor in enumerate(images):\n",
    "        img = tensor_to_pil(img_tensor)\n",
    "        plt.subplot(1, len(images), i+1)\n",
    "        plt.imshow(img)\n",
    "        plt.axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def tensor_to_pil(img):\n",
    "    img = (img / 2 + 0.5).clamp(0, 1)\n",
    "    img = img.cpu().float().permute(1, 2, 0).numpy()\n",
    "    return Image.fromarray((img * 255).astype(\"uint8\"))\n",
    "\n",
    "def process_latent_file(file_path, vae):\n",
    "    print(f\"Processing: {file_path}\")\n",
    "    latent = load_latent(file_path)\n",
    "    if latent is not None:\n",
    "        decoded_images = decode_latents(latent.unsqueeze(0), vae)\n",
    "        show_images(decoded_images, name=f\"Decoded: {os.path.basename(file_path)}\")\n",
    "\n",
    "def process_directory(directory, vae):\n",
    "    for file in os.listdir(directory):\n",
    "        if file.endswith('.safetensors'):\n",
    "            process_latent_file(os.path.join(directory, file), vae)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    vae = load_vae(device)\n",
    "    \n",
    "    latent_file = r\"G:\\comfy_dev\\ComfyUI\\input\\img_001_MxN4PuRrT4mg5FmEPGDf8g.safetensors\"\n",
    "    process_latent_file(latent_file, vae)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from diffusers import AutoencoderTiny\n",
    "import safetensors\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def load_vae(device='cuda' if torch.cuda.is_available() else 'cpu'):\n",
    "    vae = AutoencoderTiny.from_pretrained(\"madebyollin/taef1\", torch_dtype=torch.bfloat16).to(device)\n",
    "    return vae\n",
    "\n",
    "def load_latent(path):\n",
    "    try:\n",
    "        with safetensors.safe_open(path, framework=\"pt\", device=\"cpu\") as f:\n",
    "            latent = f.get_tensor('latent')\n",
    "        return latent\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading latent from {path}: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "def decode_latents(latents: torch.Tensor, vae: AutoencoderTiny):\n",
    "    latents = latents.to(vae.device, dtype=torch.bfloat16)\n",
    "    latents = latents / vae.config.scaling_factor\n",
    "    with torch.no_grad():\n",
    "        images = vae.decode(latents).sample\n",
    "    return images\n",
    "\n",
    "def tensor_to_pil(img):\n",
    "    img = (img / 2 + 0.5).clamp(0, 1)\n",
    "    img = img.cpu().float().permute(1, 2, 0).numpy()\n",
    "    return Image.fromarray((img * 255).astype(\"uint8\"))\n",
    "\n",
    "def process_latent_file(file_path, vae, output_dir):\n",
    "    print(f\"Processing: {file_path}\")\n",
    "    latent = load_latent(file_path)\n",
    "    if latent is not None:\n",
    "        decoded_images = decode_latents(latent.unsqueeze(0), vae)\n",
    "        for i, img_tensor in enumerate(decoded_images):\n",
    "            img = tensor_to_pil(img_tensor)\n",
    "            output_filename = os.path.join(output_dir, f\"{os.path.splitext(os.path.basename(file_path))[0]}_{i}.png\")\n",
    "            img.save(output_filename)\n",
    "            print(f\"Saved: {output_filename}\")\n",
    "\n",
    "def process_directory(input_dir, vae):\n",
    "    output_dir = os.path.join(input_dir, \"decoded\")\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    \n",
    "    for file in os.listdir(input_dir):\n",
    "        if file.endswith('.safetensors'):\n",
    "            process_latent_file(os.path.join(input_dir, file), vae, output_dir)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    vae = load_vae(device)\n",
    "    \n",
    "    input_dir = r\"P:\\downloads\\latents (5)\\imagecache\"\n",
    "    process_directory(input_dir, vae)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from PIL import Image\n",
    "import sys\n",
    "from PIL.ExifTags import TAGS\n",
    "from IPython.display import display, Image as IPImage\n",
    "\n",
    "def read_png_metadata(file_path):\n",
    "    try:\n",
    "        with Image.open(file_path) as img:\n",
    "            # Display the image\n",
    "            display(IPImage(filename=file_path))\n",
    "            \n",
    "            # Check if it's a PNG file\n",
    "            if img.format != 'PNG':\n",
    "                print(f\"Error: {file_path} is not a PNG file.\")\n",
    "                return\n",
    "\n",
    "            print(f\"Metadata for {os.path.basename(file_path)}:\")\n",
    "            \n",
    "            # Basic image information\n",
    "            print(f\"  Format: {img.format}\")\n",
    "            print(f\"  Mode: {img.mode}\")\n",
    "            print(f\"  Size: {img.size}\")\n",
    "            \n",
    "            # Display all EXIF data\n",
    "            exif_data = img._getexif()\n",
    "            if exif_data:\n",
    "                print(\"  EXIF Data:\")\n",
    "                for tag_id, value in exif_data.items():\n",
    "                    tag = TAGS.get(tag_id, tag_id)\n",
    "                    print(f\"    {tag}: {value}\")\n",
    "            else:\n",
    "                print(\"  No EXIF data found\")\n",
    "            \n",
    "            # Other metadata\n",
    "            print(\"  Other metadata:\")\n",
    "            for key, value in img.info.items():\n",
    "                if key != 'exif':  # Skip exif data as we've already displayed it\n",
    "                    print(f\"    {key}: {value}\")\n",
    "\n",
    "            # If there's no additional metadata\n",
    "            if not img.info and not exif_data:\n",
    "                print(\"    No additional metadata found.\")\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error reading {file_path}: {str(e)}\")\n",
    "\n",
    "def process_directory(directory):\n",
    "    for filename in os.listdir(directory):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            file_path = os.path.join(directory, filename)\n",
    "            read_png_metadata(file_path)\n",
    "            print()  # Empty line for separation between files\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    path = r\"P:\\datasets\\Posters\\base\\img_013.png\"\n",
    "    \n",
    "    if os.path.isfile(path):\n",
    "        read_png_metadata(path)\n",
    "    elif os.path.isdir(path):\n",
    "        process_directory(path)\n",
    "    else:\n",
    "        print(f\"Error: {path} is not a valid file or directory.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from diffusers import AutoencoderTiny\n",
    "import safetensors\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def load_vae(device='cuda' if torch.cuda.is_available() else 'cpu'):\n",
    "    vae = AutoencoderTiny.from_pretrained(\"madebyollin/taef1\", torch_dtype=torch.bfloat16).to(device)\n",
    "    return vae\n",
    "\n",
    "def load_latent(path):\n",
    "    try:\n",
    "        if path.endswith('.safetensors'):\n",
    "            with safetensors.safe_open(path, framework=\"pt\", device=\"cpu\") as f:\n",
    "                latent = f.get_tensor('latent')\n",
    "        elif path.endswith('.pt'):\n",
    "            latent = torch.load(path, map_location='cpu')\n",
    "            if isinstance(latent, dict) and 'latent' in latent:\n",
    "                latent = latent['latent']\n",
    "        else:\n",
    "            raise ValueError(f\"Unsupported file format: {path}\")\n",
    "        return latent\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading latent from {path}: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "def decode_latents(latents: torch.Tensor, vae: AutoencoderTiny):\n",
    "    latents = latents.to(vae.device, dtype=torch.bfloat16)\n",
    "    latents = latents / vae.config.scaling_factor\n",
    "    with torch.no_grad():\n",
    "        images = vae.decode(latents).sample\n",
    "    return images\n",
    "\n",
    "def tensor_to_pil(img):\n",
    "    img = (img / 2 + 0.5).clamp(0, 1)\n",
    "    img = img.cpu().float().permute(1, 2, 0).numpy()\n",
    "    return Image.fromarray((img * 255).astype(\"uint8\"))\n",
    "\n",
    "def process_latent_file(file_path, vae, output_dir):\n",
    "    print(f\"Processing: {file_path}\")\n",
    "    latent = load_latent(file_path)\n",
    "    if latent is not None:\n",
    "        decoded_images = decode_latents(latent.unsqueeze(0), vae)\n",
    "        for i, img_tensor in enumerate(decoded_images):\n",
    "            img = tensor_to_pil(img_tensor)\n",
    "            output_filename = os.path.join(output_dir, f\"{os.path.splitext(os.path.basename(file_path))[0]}_{i}.png\")\n",
    "            img.save(output_filename)\n",
    "            print(f\"Saved: {output_filename}\")\n",
    "\n",
    "def process_directory(input_dir, vae):\n",
    "    output_dir = os.path.join(input_dir, \"decoded\")\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    \n",
    "    for file in os.listdir(input_dir):\n",
    "        if file.endswith(('.safetensors', '.pt')):\n",
    "            process_latent_file(os.path.join(input_dir, file), vae, output_dir)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    vae = load_vae(device)\n",
    "    \n",
    "    input_dir = r\"P:\\downloads\\latents (5)\\imagecache\"\n",
    "    process_directory(input_dir, vae)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: diffusers\n",
      "Version: 0.25.0\n",
      "Summary: State-of-the-art diffusion in PyTorch and JAX.\n",
      "Home-page: https://github.com/huggingface/diffusers\n",
      "Author: The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)\n",
      "Author-email: patrick@huggingface.co\n",
      "License: Apache 2.0 License\n",
      "Location: g:\\anaconda3\\envs\\aes\\lib\\site-packages\n",
      "Requires: filelock, huggingface-hub, importlib-metadata, numpy, Pillow, regex, requests, safetensors\n",
      "Required-by: lycoris_lora\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Ignoring invalid distribution -orch (c:\\users\\daniel\\appdata\\roaming\\python\\python310\\site-packages)\n",
      "WARNING: Ignoring invalid distribution - (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -mpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -mpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -ympy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: accelerate\n",
      "Version: 0.25.0\n",
      "Summary: Accelerate\n",
      "Home-page: https://github.com/huggingface/accelerate\n",
      "Author: The HuggingFace team\n",
      "Author-email: sylvain@huggingface.co\n",
      "License: Apache\n",
      "Location: g:\\anaconda3\\envs\\aes\\lib\\site-packages\n",
      "Requires: huggingface-hub, numpy, packaging, psutil, pyyaml, safetensors, torch\n",
      "Required-by: peft\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Ignoring invalid distribution -orch (c:\\users\\daniel\\appdata\\roaming\\python\\python310\\site-packages)\n",
      "WARNING: Ignoring invalid distribution - (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -mpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -mpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -umpy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -ympy (g:\\anaconda3\\envs\\aes\\lib\\site-packages)\n"
     ]
    }
   ],
   "source": [
    "!pip show diffusers  \n",
    "!pip show accelerate"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
